#include "cc_math_tool.h"
#include "cnpy.h"
#include "process_boxes_with_nms.h"
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iostream>
#include <string.h>

#include <opencv2/opencv.hpp>

// 定义类别颜色映射 (BGR格式)
std::vector<cv::Scalar> class_colors = {
    cv::Scalar(0, 0, 255),   // 红
    cv::Scalar(0, 255, 0),   // 绿
    cv::Scalar(255, 0, 0),   // 蓝
    cv::Scalar(0, 255, 255), // 黄
    cv::Scalar(255, 0, 255), // 紫
    cv::Scalar(255, 255, 0)  // 青
};

template <typename _Tp>
int activation_function_softmax(const _Tp *src, _Tp *dst, int length) {
  const _Tp alpha = *std::max_element(src, src + length);
  _Tp denominator{0};
  for (int i = 0; i < length; ++i) {
    dst[i] = std::exp(src[i] - alpha);
    denominator += dst[i];
  }

  for (int i = 0; i < length; ++i) {
    dst[i] /= denominator;
  }
  return 0;
}

void decode(float *conv[6]) {
  int inpHeight_ = 640;
  int inpWidth_ = 960;
  float confThreshold_ = 0.5;
  float nmsThreshold_ = 0.5;
  int classNum_ = 6;

  float *bbox_conv[3]; //框
  float *cube_conv[3]; //分数
  for (int i = 0; i < 3; i++) {
    cube_conv[i] = conv[i];
    bbox_conv[i] = conv[i + 3];
    printf("bbox_conv:%p, cube_conv:%p\n", bbox_conv[i], cube_conv[i]);
  }
  //   std::vector<ObjCubeBBox> objBBoxs;
  std::vector<CcObjBBox> objBBoxs_nms;

  std::vector<int> sampling_rate = {8, 16, 32};
  std::vector<float> narray = {0.0f,  1.0f,  2.0f,  3.0f, 4.0f,  5.0f,
                               6.0f,  7.0f,  8.0f,  9.0f, 10.0f, 11.0f,
                               12.0f, 13.0f, 14.0f, 15.0f};

  for (int i = 0; i < 3; i++) {
    std::cout << "sampling_rate: " << sampling_rate[i] << std::endl;
    int h = inpHeight_ / sampling_rate[i];
    int w = inpWidth_ / sampling_rate[i];
    printf("h:%d, w:%d\n", h, w);

    int feature_strip = w * h;
    auto box_feat = bbox_conv[i];
    auto score_feat = cube_conv[i];
    std::vector<int> x_vec;
    std::vector<int> y_vec;
    std::vector<float> max_score;
    std::vector<int> max_class;

    std::vector<float> best_boxs;

    std::vector<float> best_scores;
    //先分数

    max_class.clear();
    int lenth = 0;
    for (int j = 0; j < feature_strip; j++) {
      int x = j / w;
      int y = j % w;

      best_scores.clear();
      for (int c = 0; c < classNum_; c++) {
        best_scores.push_back(*(score_feat + j + (c * feature_strip)));
      }

      auto max_iter = std::max_element(best_scores.begin(), best_scores.end());
      float max_value = *max_iter; // 解引用获取值
      size_t max_index =
          std::distance(best_scores.begin(), max_iter); // 获取索引

      float score = sigmoid(max_value);
      if (score <= confThreshold_)
        continue;

      std::cout << "score: " << score << std::endl;

      x_vec.push_back(x);
      y_vec.push_back(y);
      max_score.emplace_back(score);
      max_class.emplace_back(max_index);

      float *data_ptr = box_feat;
      for (int k = 0; k < 64; k++) {
        best_boxs.emplace_back(*(data_ptr + j + (k * feature_strip)));
      }
    }

    if (max_score.size() == 0) {
      std::cout << __FILE__ << ":" << __LINE__ << " max_score size "
                << max_score.size() << std::endl;

      continue;
    }

    std::vector<float> final_boxs;
    int len = 0;
    for (int m = 0; m < max_score.size() * 4; m++) {
      std::vector<float> temp_boxs;
      std::vector<float> temp;

      temp_boxs.resize(16);
      temp.resize(16);
      memcpy(temp_boxs.data(), best_boxs.data() + len, sizeof(float) * 16);
      len += 16;

      float sum = 0.0;
      activation_function_softmax(temp_boxs.data(), temp.data(), 16);
      for (int k = 0; k < 16; k++) {
        sum += temp[k] * narray[k];
      }
      // printf("sum:%f\n", sum);
      final_boxs.emplace_back(sum);
    }

    int kps_len = 0;

    std::cout << __FILE__ << ":" << __LINE__ << " AFTER max_score size "
              << max_score.size() << std::endl;
    for (int k = 0; k < max_score.size(); k++) {
      int hh = x_vec[k];
      int ww = y_vec[k];

      float score = max_score[k];
      //根据分数置信度过滤
      // if (score <= confThreshold_)
      //     continue;

      CcObjBBox objbbox;

      float x0 = final_boxs[k * 4 + 0];
      float y0 = final_boxs[k * 4 + 1];
      float x1 = final_boxs[k * 4 + 2];
      float y1 = final_boxs[k * 4 + 3];

      x0 = (ww + 0.5 - x0) * sampling_rate[i];
      y0 = (hh + 0.5 - y0) * sampling_rate[i];
      x1 = (ww + 0.5 + x1) * sampling_rate[i];
      y1 = (hh + 0.5 + y1) * sampling_rate[i];

      objbbox.bbox.x = x0;
      objbbox.bbox.y = y0;
      objbbox.bbox.width = x1 - x0;
      objbbox.bbox.height = y1 - y0;
      objbbox.score = score;
      objbbox.label = max_class[k];

      objBBoxs_nms.push_back(objbbox);
    }
  }
  printf("objBBoxs_nms:%d\n", objBBoxs_nms.size());
  std::vector<int> out_nms = nms(objBBoxs_nms, nmsThreshold_, classNum_);
  printf("out_nms:%d\n", out_nms.size());
  std::vector<CcObjBBox> outBBoxs;
  for (int i : out_nms) {
    outBBoxs.push_back(objBBoxs_nms[i]);
  }

  cv::Mat org_show = cv::imread("test_car2.jpg");
  cv::resize(org_show, org_show, cv::Size(960, 640));
  for (auto k : outBBoxs) {
    // 根据类别选择颜色
    cv::Scalar color = class_colors[k.label % class_colors.size()];
    cv::rectangle(org_show, cv::Point(k.bbox.x, k.bbox.y),
                  cv::Point(k.bbox.x + k.bbox.width, k.bbox.y + k.bbox.height),
                  color, 2);

    // 构建标签文本
    std::string label = cv::format("Class%d: %.2f", k.label, k.score);

    // 计算文本位置
    int baseline = 0;
    cv::Size text_size =
        cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);

    // 画文本背景
    cv::rectangle(
        org_show, cv::Point(k.bbox.x, k.bbox.y - text_size.height - 5),
        cv::Point(k.bbox.x + text_size.width, k.bbox.y), color, cv::FILLED);

    // 写文本
    cv::putText(org_show, label, cv::Point(k.bbox.x, k.bbox.y - 5),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
  }
  cv::imwrite("result.jpg", org_show);

  //   return outBBoxs;
  return;
}

void process_boxes(float *score[3], float *box[3], const int h[3],
                   const int w[3], float scale_x, float scale_y,
                   float conf_thresh, int reg_max,
                   std::vector<std::vector<float>> &out_boxes,
                   std::vector<float> &out_scores) {
  out_boxes.clear();
  out_scores.clear();

  for (int l = 0; l < 3; ++l) {
    int H = h[l], W = w[l];
    int stride = 8 << l;               // 8, 16, 32
    const float *score_ptr = score[l]; // [1, 6, H, W]
    const float *box_ptr = box[l];     // [1, 64, H, W]
    int score_area = H * W;
    int box_area = H * W;
    int n_classes = 6;

    std::vector<float> prob(reg_max);

    for (int y = 0; y < H; ++y) {
      for (int x = 0; x < W; ++x) {
        // 找最大类别分数
        int offset = y * W + x;
        float max_score = -1e10f;
        for (int c = 0; c < n_classes; ++c) {
          float cls_score = score_ptr[c * score_area + offset];
          if (cls_score > max_score)
            max_score = cls_score;
        }
        max_score = sigmoid(max_score);
        if (max_score < conf_thresh)
          continue;

        // box_weights: [4]
        float box_weights[4] = {0};
        for (int b = 0; b < 4; ++b) {
          // 每个方向16个logit
          float logits[16];
          for (int k = 0; k < reg_max; ++k) {
            logits[k] = box_ptr[(b * reg_max + k) * box_area + offset];
          }
          softmax(logits, reg_max, prob.data());
          for (int k = 0; k < reg_max; ++k) {
            box_weights[b] += prob[k] * k;
          }
        }
        float x0 = box_weights[0];
        float y0 = box_weights[1];
        float x1 = box_weights[2];
        float y1 = box_weights[3];

        std::vector<float> box_vec = {(x + 0.5f - x0) * stride * scale_x,
                                      (y + 0.5f - y0) * stride * scale_y,
                                      (x + 0.5f + x1) * stride * scale_x,
                                      (y + 0.5f + y1) * stride * scale_y};
        out_boxes.push_back(box_vec);
        out_scores.push_back(max_score);
      }
    }
  }
}

int main() {
  std::cout << "Hello World!" << std::endl;

  cnpy::NpyArray a = cnpy::npy_load("./output0/output0.npy");
  //   if (a.fortran_order) {
  //     std::cout << "a is fortran_order" << std::endl;
  //   }
  //   std::cout << a.shape[0] << std::endl;
  //   std::cout << a.shape[1] << std::endl;
  //   std::cout << a.shape[2] << std::endl;
  //   std::cout << a.shape[3] << std::endl;
  cnpy::NpyArray a1 = cnpy::npy_load("./output0/output1.npy");
  //   std::cout << a1.shape[0] << " " << a1.shape[1] << " " << a1.shape[2] << "
  //   "
  //             << a1.shape[3] << std::endl;
  cnpy::NpyArray a2 = cnpy::npy_load("./output0/output2.npy");
  //   std::cout << a2.shape[0] << " " << a2.shape[1] << " " << a2.shape[2] << "
  //   "
  //             << a2.shape[3] << std::endl;

  cnpy::NpyArray b = cnpy::npy_load("./output0/output3.npy");
  cnpy::NpyArray b1 = cnpy::npy_load("./output0/output4.npy");
  cnpy::NpyArray b2 = cnpy::npy_load("./output0/output5.npy");
  //   std::cout << b.shape[0] << std::endl;
  //   std::cout << b.shape[1] << std::endl;
  //   std::cout << b.shape[2] << std::endl;
  //   std::cout << b.shape[3] << std::endl;

  std::vector<float *> score;
  score.push_back(a.data<float>());
  score.push_back(a1.data<float>());
  score.push_back(a2.data<float>());

  std::vector<float *> box;
  box.push_back(b.data<float>());
  box.push_back(b1.data<float>());
  box.push_back(b2.data<float>());

  std::vector<float *> total;
  auto aa = a.as_vec<float>();
  total.push_back(aa.data());
  auto aa1 = a1.as_vec<float>();
  total.push_back(aa1.data());
  auto aa2 = a2.as_vec<float>();
  total.push_back(aa2.data());
  auto bb = b.as_vec<float>();
  total.push_back(bb.data());
  auto bb1 = b1.as_vec<float>();
  total.push_back(bb1.data());
  auto bb2 = b2.as_vec<float>();
  total.push_back(bb2.data());

#if 0
  std::ofstream out("output.txt");
  out << std::fixed << std::setprecision(6); // 固定6位小数
  for (const auto &val : aa1) {
    out << val << "\n"; // 每行写入一个浮点数
  }
  out.close();
#endif
  //   cnpy::NpyArray a = cnpy::npy_load("output0.npy");
  //   float *ptr = a.data<float>();
  //   size_t total_elements = 1;
  //   for (auto dim : a.shape)
  //     total_elements *= dim;

  //   // 检查指针偏移量连续性
  //   for (size_t i = 1; i < total_elements; ++i) {
  //     assert(&ptr[i] - &ptr[i - 1] == 1); // 连续内存应满足相邻元素地址差=1
  //   }

  //   float *ptr = a1.data<float>();
  //   for (size_t i = 0; i < a1.shape[0]; i++) {
  //     for (size_t j = 0; j < a1.shape[1]; j++) {
  //       float val = ptr[i * a1.shape[1] + j]; // 行优先计算索引
  //       if (val > 1.0f)
  //         std::cout << val << std::endl;
  //     }
  //   }

  //   std::cout << "total size: " << total.size() << std::endl;
  //   std::cout << "\nword size: " << a1.word_size << std::endl;
  //   if (a1.word_size == sizeof(float)) {
  //     std::cout << "float" << std::endl;
  //   }

#if 1
  decode(total.data());
#else
  int h[] = {80, 40, 20};
  int w[] = {120, 60, 30};
  //   std::vector<std::vector<float>> out_boxes;
  //   std::vector<float> out_scores;
  //   process_boxes(score.data(), box.data(), h, w, 1.0f, 1.0f, 0.5, 16,
  //   out_boxes,
  //                 out_scores);

  std::vector<Detection> out_dets;
  process_boxes_with_nms(score.data(), box.data(), h, w, 1.0f, 1.0f, 0.5, 16,
                         0.5, 10, out_dets);

  std::cout << "out_dets size: " << out_dets.size() << std::endl;
  cv::Mat org_show = cv::imread("test_car.jpg");
  cv::resize(org_show, org_show, cv::Size(960, 640));
  for (auto k : out_dets) {
    // 根据类别选择颜色
    cv::Scalar color = class_colors[k.class_id % class_colors.size()];

    cv::rectangle(org_show, cv::Point(k.box[0], k.box[1]),
                  cv::Point(k.box[2], k.box[3]), color, 2);

    // 构建标签文本
    std::string label = cv::format("Class%d: %.2f", k.class_id, k.score);

    // 计算文本位置
    int baseline = 0;
    cv::Size text_size =
        cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseline);

    // 画文本背景
    cv::rectangle(
        org_show, cv::Point(k.box[0], k.box[1] - text_size.height - 5),
        cv::Point(k.box[0] + text_size.width, k.box[1]), color, cv::FILLED);

    // 写文本
    cv::putText(org_show, label, cv::Point(k.box[0], k.box[1] - 5),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
  }
  cv::imwrite("result.jpg", org_show);
#endif
  return 0;
}